'''
基于模式分解算法模拟部分相干光，生成gsm光场，光场传播
'''
import copy
from re import X
from statistics import mode
import time
from datetime import datetime

import cv2
import numpy as np
from scipy import constants as c
from scipy import signal, special
from tqdm.auto import tqdm

exp = np.exp
fft = np.fft
Hermite = special.eval_hermite   # 厄密多项式
factorial = special.factorial   # 阶乘,ny.math 中的阶乘函数不支持数组变量
cztsp = signal.czt

class OpticalField(object):
    def __init__(self,mode,occup,x,y,k):
        '''
        Description
        -----------
        部分相干光互强度的模式分解形式

        Args
        ----
            mode (np.array): 相干模式 (nmx*nmy*nx*ny);
                `nmx(y)`: x(y) 方向模式数量; `nx(y)`: x(y) 方向采样网格点数
            occup (tuple(np.array,np.array) ): 模式占有率 (occupancy), tuple(array(nmx,1,1,1), array(1,nmy,1,1))
            x (np.array): x 坐标 (nx,)
            y (np.array): y 坐标 (ny,)
        '''
        pass
        self.mode = mode
        self.occup =occup
        self.x = x
        self.y = y
        self.k = k
    
    def calcCoherence(self,angle=0,bias=0,threshold=1e3):
        '''
        Description
        -----------
        使用相干模式(CM)计算相干信息: 互强度(mulInt)、光强(I) 和复相干度(cohdeg)

        Args
        ----
            angle (int, optional): 方向角[deg]，沿直线 y=tan(angle)*x + bias 计算相干信息. Defaults to 0.
            bias (int, optional): 偏移[像素]. Defaults to 0.
            threshold (float, optional): 计算相干度时的滤波阈值. Defaults to 0.

        Returns
        -------
            coh (dict): coh:{mulInt: 互强度, I: 光强, cohdeg: 复相干度}
        '''
        occup_x,occup_y = self.occup
        En = self.mode
        # nmx(y):occup_x(y) 模式数量，nx(y): x(y) 方向采样网格点数
        nmx,nmy,nx,ny = En.shape
        
        # 旋转 En
        if angle%90 == 0:
            En_r = np.rot90(En,-angle//90,axes=(-2,-1))
        else:
            center = (nx // 2, ny // 2)
            M = cv2.getRotationMatrix2D(center, angle, 1)
            li = []
            for ii in range(nmx):
                lj = []
                for jj in range(nmy):
                    tmp = cv2.warpAffine(np.real(En[ii,jj]), M, (nx,ny)) \
                                + 1j*cv2.warpAffine(np.imag(En[ii,jj]), M, (nx,ny))
                    lj.append(tmp)
                li.append(lj)
            En_r = np.array(li)
        # 构建每个模式的互强度
        En1 = En_r[:,:,:,nx//2+bias]
        En12 = En1[:,:,None,:] * En1[:,:,:,None].conjugate()

        # 计算整个波前的互强度
        mulInt = 0* np.zeros([nx,nx],
            dtype=np.complex128)
        for k1,v1 in enumerate(occup_x):
            for k2,v2 in enumerate(occup_y):
                mulInt += (v1*v2) * En12[k1,k2]
        # 光强
        I = mulInt.diagonal()
        
        # 谱复干度
        np.seterr(invalid='ignore')     # 忽略下面除法的 NaN 或 inf
        cohdeg = mulInt / (np.outer(I,I)**.5)
        # 将 inf 的区域置零
        cohdeg[cohdeg==np.inf] = 0
        # 将光强低于阈值的区域置零
        zone = abs(I)>abs(I).max()/threshold
        zone = zone * zone[:,None]
        cohdeg[~zone] = 0
        
        coh = {'mulInt':mulInt, 'I':I, 'cohdeg':cohdeg}
        return coh
    
    def qFunc(self,x,y,z):
        '''菲涅尔近似下,传播积分等价于源场与一个二次位相因子乘积后的F.T.，这里定义的 q 函数用于计算这个二次位相因子

        Parameters
        ----------
        x :
            横坐标
        y : [type]
            纵坐标
        z : [type]
            轴向传播距离

        Returns
        -------
        [type]
            菲涅尔近似下的二次位相弯曲
        '''        
        q = np.exp(1j*self.k/2/z
                   *(x.reshape(-1,1)**2+y**2))
        return q
    
    def propagate(self,z,coord=None,T=None):
        '''计算场的传播

        Parameters
        ----------
        z : [float]
            传播距离
        coord : list, optional
            [[lx,rx,nx],[ly,ry,ny]], l:left,r:right,n:number;by default None
        T : string, optional
            optical element name, by default None
        '''
        Es = copy.deepcopy(self.mode)
        # 源场坐标
        xs  = self.x
        ys  = self.y
        dxs = xs[1]-xs[0]; dys = ys[1]-ys[0]
        # 离散 FT 对应的传播场坐标 range
        x_range = 2*c.pi/self.k*z / dxs
        y_range = 2*c.pi/self.k*z / dys
        # 生成传播场坐标
        if coord is not None:
            # 如果有指定则使用指定的坐标
            xp0,xp1,nx = coord[0]; yp0,yp1,ny = coord[1]
        else:
            # 没有指定则使用 ft 允许的最大坐标
            xp0 = x_range/2; xp1 = -xs; nx = xs.size
            yp0 = y_range/2; yp1 = -ys; ny = ys.size
        xp = np.linspace(xp0,xp1,nx)
        yp = np.linspace(yp0,yp1,ny)
        
        # 屏函数
        if T is None:
            pass
        else:
            Es *= T[None,None,:,:]
        
        # 初始化传播场
        Ep = np.empty(Es.shape[:2]+(nx,ny),dtype=np.complex128)
        
        a=datetime.now()
        # 源场 q 因子
        q_s = self.qFunc(xs,ys,z)
        Es *= q_s.reshape((1,1)+q_s.shape)
        
        # 创建进度条
        loop_len = Es[:,:,0,0].size
        pbar = tqdm(total=loop_len)
        # 执行 F.T. 操作
        range_ = [(xp1-xp0)/x_range,(yp1-yp0)/y_range]
        for k1,v1 in enumerate(Es):
            for k2,v2 in enumerate(v1):
                Ep[k1,k2] = ft2_czt(v2,
                    range_=range_,mesh=[nx,ny])
                pbar.update()
        pbar.close()    # 关闭进度条
        
        # 传播场 q 因子
        q_p = self.qFunc(xp,yp,z)
        Ep *= dxs*dys * exp(1j*self.k*z) \
            * self.k/ (2*c.pi*1j) / z \
            * q_p.reshape((1,1)+q_p.shape)
        b=datetime.now()
        print(f'计算传播耗时：{(b-a).seconds} s')
        # 返回传播场
        return OpticalField(Ep,xp,yp,self.occup,self.k)

    def intensity(self):
        '''计算 2D 光强'''
        occup_x = self.occup[0][:,None,None,None]
        occup_y = self.occup[1][None,:,None,None]
        return np.sum( 
            abs(self.mode)**2 * (occup_x*occup_y),
            axis=(0, 1) )

    def wfAngle(self):
        '''计算关于坐标原点的 2D 相位'''
        # 模式占有率
        occup_x = self.occup[0][:,None,None,None]
        occup_y = self.occup[1][None,:,None,None]
        _,_,nx,ny = self.mode.shape
        cx,cy = nx//2,ny//2
        # 不同模式在坐标原点的取值
        E_center = self.mode[:,:,cx,cy][:,:,None,None]
        return np.angle(np.sum(
            E_center.conj() * self.mode * (occup_x*occup_y),
            axis=(0, 1) ))


class GSMField(OpticalField):
    def __init__(self,sig_i,sig_c,k,s0,x,y,z,nmx,nmy):
        '''
        Description
        -----------
        计算 GSM 的相干模式光场

        Args
        ----
            `sig_i ([float,float])`: 光强的 RMS [sig_ix,sig_iy]
            `sig_c ([float,float])`: 相干函数的 RMS [sig_cx,sig_cy]
            `k (float)`: 波束 k
            `s0 (float)`: 光强系数
            `x (np.array)`: x 坐标
            `y (np.array)`: y 左边
            `z (float)`: 计算距离束腰 z 处的光场
            `nmx (int)`: x 方向模式数量
            `nmy (int)`: y 方向模式数量
        '''
        self.sig_i = np.array(sig_i)
        self.sig_c = np.array(sig_c)
        self.k = k
        self.s0 = s0
        self.x = x
        self.y = y
        self.z = z
        self.nmx = nmx
        self.nmy = nmy
        self.calcMode()
        self.calcOccup()
        super(GSMField,self).__init__(self.mode,self.occup,x,y,k)
        
    def gsmWaist(self):
        r'''
        Description
        -----------
        给出波矢为 k 的 GSM 光在 z 处的参数, 计算束腰参数.
        
        Markdown
        --------
        波矢为 $k$ 的 GSM 光, 给出 z 处光斑 RMS $\Sigma$ 和相干 RMS $\Xi$, 求束腰处相应的 $\sigma$ 和 $\xi$
        $$\sigma = \frac{1}{\sqrt{2}}\left[\Sigma^2 - \sqrt{\Sigma^4 - 4\left(\frac{2z}{\zeta k}\right)^2} ~ \right]^{\frac{1}{2}} \\
        \xi = \Xi * \frac{\sigma}{\Sigma}$$
        式中用到的 $\zeta$ 通过下式求出:
        $$\zeta = \left[  1+(\frac{2\Sigma}{\Xi})^2 \right]^{\frac{1}{2}}$$
        '''
        sig_iz = self.sig_i
        sig_cz = self.sig_c
        # 辅助参数
        self.zeta = 1/np.sqrt(1+(2*sig_iz/sig_cz)**2)
        # 计算束腰处光强 (相干) 分布的 RMS
        self.sig_iw = (sig_iz**2 + (sig_iz**4 - 4*(self.z/self.k/self.zeta)**2 )**.5 )**.5 /2**.5
        self.sig_cw = sig_cz * self.sig_iw/sig_iz
        # 有效距离
        self.z_eff = 2*self.k*self.sig_iw**2*self.zeta
    
    def calcMode(self):
        '''
        Description
        -----------
        计算 GSM 在 z 处的相干模式
        '''
        self.gsmWaist()     # 计算束腰参数
        # 坐标系和波数 k
        z = self.z
        x = self.x[None,None,:,None]
        y = self.y[None,None,None,:]
        k = self.k
        
        # 辅助参数 zeta 和有效距离 z_eff
        zeta_x  = self.zeta[0] ; zeta_y  = self.zeta[1]
        z_eff_x = self.z_eff[0]; z_eff_y = self.z_eff[1]
        # z 处波前曲率半径和光斑 RMS
        self.R_z  = (z+1e-50)*(1+(self.z_eff/(z+1e-50))**2)
        R_zx   = self.R_z[0];   R_zy   = self.R_z[1]
        sig_ix = self.sig_i[0]; sig_iy = self.sig_i[1]
        
        # j_m(n)：x(y) 方向相干模式的阶次
        j_m = np.linspace(0,self.nmx,self.nmx+1)[:,None,None,None]
        j_n = np.linspace(0,self.nmy,self.nmy+1)[None,:,None,None]
        # 计算 z 处的相干模式
        En_x = ( (2*c.pi*sig_ix**2*zeta_x)**-.25
            * (2**j_m * factorial(j_m))**-.5
            * Hermite(np.int32(j_m), x/ (sig_ix*(2*zeta_x)**.5))
            * exp(-x**2 / (4* sig_ix**2 *zeta_x))
            * exp(1j*(k*z- (j_m+1)*np.math.atan(z/z_eff_x) )
                +1j*k* x**2 / (2*R_zx)) )
        En_y = ( (2*c.pi*sig_iy**2*zeta_y)**-.25
            * (2**j_n * factorial(j_n))**-.5
            * Hermite(np.int32(j_n), y/ (sig_iy*(2*zeta_y)**.5))
            * exp(-y**2 / (4* sig_iy**2 *zeta_y))
            * exp(1j*(k*z- (j_n+1)*np.math.atan(z/z_eff_y) )
                +1j*k* y**2 / (2*R_zy)) )
        self.mode = En_x * En_y
    
    def calcOccup(self):
        '''
        Description
        -----------
        计算模式占有率
        '''
        zeta  = self.zeta
        m,n = self.nmx, self.nmy
        j_m = np.linspace(0,m,m+1)[:,None]
        j_n = np.linspace(0,n,n+1)[:,None]
        # 计算光源(z=0)处每个模式的系数
        occup0 = ( (2*c.pi*self.sig_iw**2
            *self.s0)**.5 * 2*zeta / (1+zeta) )
        kappa = ((1-zeta)/(1+zeta))
        occup_x = (occup0 * kappa**j_m)[:,0]
        occup_y = (occup0 * kappa**j_n)[:,1]
        self.occup = (occup_x,occup_y)


def elec2optic(emittance,beta,phoEng,L):
    '''电子束参数 -> GSM 参数

    Parameters
    ----------
    emittance : [type]
        电子束发射度
    beta : [type]
        电子束流 beta 函数
    phoEng : [type]
        光子能量
    L : [type]
        插入件长度

    Returns
    -------
    [type]
        GSM 光斑半径，相干半径
    '''
    # 电子束流参数计算 GSM 参数
    wavelength = c.h*c.c/phoEng / c.e
    # 电子束流参数
    sigma_eb = (emittance*beta)**.5
    angle_eb = (emittance/beta)**.5
    # 单电子参数
    sigma_n = (2*wavelength*L)**.5 /4/np.pi
    angle_n = (wavelength/2/L)**.5
    # 光束参数
    sigma_ob = (sigma_eb**2+sigma_n**2)**.5
    angle_ob = (angle_eb**2+angle_n**2)**.5
    # 相干度
    zeta = sigma_n*angle_n / (sigma_ob*angle_ob)
    # 相干长度
    xi = 2*sigma_ob*zeta/(1-zeta**2)**.5
    return sigma_ob, xi

def ft2_czt(z,range_=[1.,1.],mesh=None,method='scipy_zoom_fft'):
    '''基于 czt 的 2D FT 变换, 可指定采样率以及频域范围
    
    Parameters
    ----------
    z : 2D array
        输入数据
    range_ : list, optional
        指定频域范围 axis[0,1], 应为 0~1 之间, 1 对应 FT 变换的默认频域范围, 取值越小频域范围越小, by default [None,None]
    mesh : list, optional
        频域采样点数量 axis[0,1], 默认与 FT 相同, 即变换前, by default [None,None]
        method: string, optional
    计算 czt 的方法: 
        1. scipy_zoom_fft 是一个特别版的 czt, 其基矢为单位圆上的等间隔频率点, 精度优于通用版 czt;
        2. scipy_czt: scipy 的 czt 函数;
    '''
    jpi = 1j*np.pi
    m0,m1 = z.shape   # 输入数据的网格
    range0,range1 = range_
    n0, n1 = mesh if mesh != None else (m0,m1)
    
    f0 = (1-range0/2) + range0/(2*n0)
    f1 = (1-range1/2) + range1/(2*n1)
    # 变换的步长
    W0 = np.exp(-2 * jpi / n0 * range0)
    W1 = np.exp(-2 * jpi / n1 * range1)
    # 变换的起点
    A_0 = np.exp(2 * jpi * f0)
    A_1 = np.exp(2 * jpi * f1)
    
    n0_arr = np.arange(n0)
    n1_arr = np.arange(n1)

    # 相位修正
    l0 = n0_arr/n0*range0 + f0
    phase_correct_x = ( np.exp( jpi* m0 *l0) )[:,None]
    l1 = n1_arr/n1*range1 + f1
    phase_correct_y = ( np.exp( jpi* m1 *l1) )[None,:]

    if method == 'scipy_zoom_fft':
        FT = signal.zoom_fft(signal.zoom_fft(
                z, [-range0,range0],n0,endpoint=True, axis=0 ),
                   [-range1,range1],n1,endpoint=True, axis=1 ) \
            * phase_correct_y*phase_correct_x
    elif method == 'scipy_czt':
        FT = cztsp(
            cztsp(z, n0, W0, A_0, axis=0),
            n1, W1, A_1, axis=1 ) * phase_correct_y*phase_correct_x
    return FT


